{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "655137e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!/usr/bin/env python\n",
    "# coding: utf-8\n",
    "# 生成HAN模型用的数据\n",
    "# In[54]:\n",
    "\n",
    "\n",
    "import pathlib\n",
    "import pickle\n",
    "\n",
    "import numpy as np\n",
    "import scipy.sparse\n",
    "import scipy.io\n",
    "import pandas as pd\n",
    "\n",
    "\n",
    "# In[55]:\n",
    "num_drug = 1482\n",
    "num_dis = 793\n",
    "num_target = 2077\n",
    "num_gene = 6265\n",
    "\n",
    "\n",
    "def get_mask():\n",
    "    dim = num_drug + num_dis\n",
    "    # print(dim)\n",
    "    type_mask = np.zeros((dim), dtype=int)\n",
    "    type_mask[num_drug:num_drug + num_dis] = 1\n",
    "    np.save('node_types_NSAP.npy',type_mask)\n",
    "    return\n",
    "\n",
    "\n",
    "def convert_data(idx_dic,type):\n",
    "    #type=0,为药物结点的邻接关系\n",
    "    if type==0:\n",
    "        dimension=num_drug\n",
    "    if type==1:\n",
    "        dimension=num_dis\n",
    "    adjm = np.zeros((dimension, dimension), dtype=int)\n",
    "    for key,vals in idx_dic.items():\n",
    "        for val in vals:\n",
    "            if type==0:\n",
    "                v1=val[-1]\n",
    "                v0=val[0]\n",
    "            if type==1:\n",
    "                v1=val[-1]-num_drug\n",
    "                v0=val[0]-num_drug\n",
    "            adjm[v1,v0]=1\n",
    "    return adjm\n",
    "\n",
    "\n",
    "def get_BioNet_data(prefix='.'):\n",
    "    in_file = open(prefix + '/0/0-1-0_idx.pickle', 'rb')\n",
    "    idx_usu = pickle.load(in_file)\n",
    "    in_file.close()\n",
    "    USU=scipy.sparse.csr_matrix(convert_data(idx_usu,0))\n",
    "\n",
    "\n",
    "    in_file = open(prefix + '/0/0-1-3-1-0_idx.pickle', 'rb')\n",
    "    idx_usgsu = pickle.load(in_file)\n",
    "    in_file.close()\n",
    "    USGSU = scipy.sparse.csr_matrix(convert_data(idx_usgsu, 0))\n",
    "\n",
    "    in_file = open(prefix + '/0/0-2-0_idx.pickle', 'rb')\n",
    "    idx_utu = pickle.load(in_file)\n",
    "    in_file.close()\n",
    "    UTU = scipy.sparse.csr_matrix(convert_data(idx_utu, 0))\n",
    "\n",
    "\n",
    "\n",
    "    in_file = open(prefix + '/1/1-0-1_idx.pickle', 'rb')\n",
    "    idx_sus = pickle.load(in_file)\n",
    "    in_file.close()\n",
    "    SUS=scipy.sparse.csr_matrix(convert_data(idx_sus,1))\n",
    "\n",
    "\n",
    "    in_file = open(prefix + '/1/1-3-1_idx.pickle', 'rb')\n",
    "    idx_sgs = pickle.load(in_file)\n",
    "    in_file.close()\n",
    "    SGS=scipy.sparse.csr_matrix(convert_data(idx_sgs,1))\n",
    "    \n",
    "    in_file = open(prefix + '/1/1-0-2-0-1_idx.pickle', 'rb')\n",
    "    idx_sgs = pickle.load(in_file)\n",
    "    in_file.close()\n",
    "    SUTUS=scipy.sparse.csr_matrix(convert_data(idx_sgs,1))\n",
    "\n",
    "    data={'USU':USU,\n",
    "          'USGSU':USGSU,\n",
    "          'UTU':UTU,\n",
    "          'SUS':SUS,\n",
    "          'SGS':SGS,\n",
    "          'SUTUS':SUTUS\n",
    "    }\n",
    "    output=open('adjData.pkl','wb')\n",
    "    pickle.dump(data,output)\n",
    "    return"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "3a1c028d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def main():\n",
    "    get_mask()\n",
    "    get_BioNet_data()\n",
    "\n",
    "\n",
    "\n",
    "if __name__=='__main__':\n",
    "     main()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41cb2220",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
